package cn.xinfei.xdecision.data.thread;

import cn.xinfei.xdecision.common.model.datax.consts.VariableConst;
import cn.xinfei.xdecision.common.model.datax.enums.ExecutorSceneEnum;
import lombok.RequiredArgsConstructor;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.MDC;
import org.springframework.core.task.TaskDecorator;

import java.util.Optional;
import java.util.UUID;

public class MdcTraceIdTaskDecorator implements TaskDecorator {


    private Boolean idenpendentTrace;

    private ExecutorSceneEnum executorEnum;

    public MdcTraceIdTaskDecorator(Boolean idenpendentTrace) {
        this.idenpendentTrace = idenpendentTrace;
    }

    public MdcTraceIdTaskDecorator(Boolean idenpendentTrace, ExecutorSceneEnum executorEnum) {
        this.idenpendentTrace = idenpendentTrace;
        this.executorEnum = executorEnum;
    }

    /**
     * 使异步线程池获得主线程的trace
     */
    @Override
    public Runnable decorate(Runnable runnable) {
        String traceId = getCurrentTraceId();
        return new MdcTraceIdRunnable(traceId, runnable);
    }

    private String getCurrentTraceId() {
        String currentTraceId;
        if (StringUtils.isNotBlank(MDC.get(VariableConst.TRACE_ID))) {
            currentTraceId = MDC.get(VariableConst.TRACE_ID);
        } else {
            currentTraceId = UUID.randomUUID().toString().replace("-", "");
        }

        if (idenpendentTrace) {
            Optional<String> executorScene = Optional.ofNullable(executorEnum)
                    .map(ExecutorSceneEnum::getCode);
            currentTraceId = currentTraceId + "_" + UUID.randomUUID().toString().replace("-", "").substring(0, 8);
            if (executorScene.isPresent()) {
                currentTraceId = currentTraceId + "_" + executorScene.get();
            }
        }
        return currentTraceId;
    }

    @RequiredArgsConstructor
    public static class MdcTraceIdRunnable implements Runnable {

        private static final String TRACE_ID_KEY = "TRACE_ID";
        private final String traceId;
        private final Runnable delegate;

        @Override
        public void run() {
            try {
                MDC.put(TRACE_ID_KEY, traceId);
                this.delegate.run();
            } finally {
                MDC.clear();
            }
        }
    }
}
